package org.spbsu.apco.server.scheduler;

import org.perf4j.StopWatch;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.spbsu.apco.common.properties.ApplicationPropertiesManager;
import org.spbsu.apco.common.properties.PropertyKey;
import org.spbsu.apco.common.task.core.BaseTask;
import org.spbsu.apco.common.task.core.BatchTask;
import org.spbsu.apco.common.task.core.Ticket;
import org.spbsu.apco.common.task.core.WaitTask;
import org.spbsu.apco.common.utils.StopWatchFactory;
import org.spbsu.apco.server.monitoring.prediction.TicketDurationPredictor;
import org.spbsu.apco.server.taskset.BaseTaskSet;
import org.springframework.beans.factory.annotation.Required;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * User: ortemij
 * Date: 05.01.13
 * Time: 22:38
 */
public abstract class BaseScheduler implements Scheduler {

    private static final Logger LOG = LoggerFactory.getLogger(BaseScheduler.class);

    private final Map<Long, Collection<Object>> finishedTaskSets = new ConcurrentHashMap<Long, Collection<Object>>();
    private final Map<Long, Object> task2result = new ConcurrentHashMap<Long, Object>();
    private final Map<Long, Long> task2taskSet = new ConcurrentHashMap<Long, Long>();
    private final Map<Long, Set<Long>> taskSet2lastTasks = new ConcurrentHashMap<Long, Set<Long>>();
    private final Map<Long, Ticket> readyForExecutionTicketsIndex = new ConcurrentHashMap<Long, Ticket>();

    private final Map<Long, Ticket> sentForExecutionTasksIndex = new HashMap<Long, Ticket>();
    private final Map<Long, BaseTask> dependingTasksIndex = new HashMap<Long, BaseTask>();
    private final Map<Long, Long> taskSet2startTime = new HashMap<Long, Long>();

    private int waitTaskDuration = Integer.valueOf(ApplicationPropertiesManager.getProperty(PropertyKey.WAIT_TASK_DURATION));
    private int batchTaskThreshold = Integer.valueOf(ApplicationPropertiesManager.getProperty(PropertyKey.BATCH_TASK_TRESHOLD));
    protected TicketDurationPredictor ticketDurationPredictor;

    @Override
    public final void holdTaskSet(BaseTaskSet taskSet) {
        LOG.debug("Scheduler started holding task set #{}", new Object[] {taskSet.getId()});
        final Map<BaseTask, List<?>> task2input = taskSet.getTask2input();

        final StopWatch stopWatch = StopWatchFactory.createOrGetLocal();
        final StopWatch stopWatchDetailed = StopWatchFactory.createOrGetLocal();
        stopWatch.start();

        // calculation of last tasks (which task set will be finished after)
        {
            final Set<BaseTask> allTasks = taskSet.getTasks();
            allTasks.removeAll(task2input.keySet());
            final Set<BaseTask> copy = new HashSet<BaseTask>(allTasks);
            for (BaseTask task : copy) {
                allTasks.removeAll(task.getSubTasks());
                dependingTasksIndex.put(task.getId(), task);
                task2taskSet.put(task.getId(), taskSet.getId());
            }
            if (!allTasks.isEmpty()) {
                taskSet2lastTasks.put(taskSet.getId(), getIds(allTasks));
            } else {
                taskSet2lastTasks.put(taskSet.getId(), getIds(task2input.keySet()));
            }
            stopWatchDetailed.lap("server__scheduler-last-task-calculation");
        }

        // tasks with no previous dependencies could be executed now
        for (Map.Entry<BaseTask, List<?>> taskAndInput : task2input.entrySet()) {
            BaseTask baseTask = taskAndInput.getKey();
            Ticket ticket = Ticket.of(baseTask, taskAndInput.getValue());
            readyForExecutionTicketsIndex.put(baseTask.getId(), ticket);
            task2taskSet.put(baseTask.getId(), taskSet.getId());
        }
        stopWatchDetailed.stop("server__scheduler-ready-tasks-preparation");

        taskSet2startTime.put(taskSet.getId(), System.currentTimeMillis());

        LOG.debug("Scheduler finished holding task set #{}", new Object[] {taskSet.getId()});
        stopWatch.stop("server__scheduler-hold-task-set");
    }

    @Override
    public final Ticket getNextTicketFor(String clientId) {
        final StopWatch stopWatch = StopWatchFactory.createOrGetLocal();

        if (readyForExecutionTicketsIndex.size() == 0) {
            reorganize();
        }

        Ticket ticket = getTicketFor(clientId);
        if (ticket == null) {
            final Ticket t = Ticket.of(new WaitTask(), Arrays.asList(waitTaskDuration));
            final long taskId = t.getTask().getId();
            sentForExecutionTasksIndex.put(taskId, t);
            readyForExecutionTicketsIndex.remove(taskId);
            stopWatch.stop("server__scheduler-next-wait-ticket");
            return t;
        }

        int duration = ticketDurationPredictor.predictDuration(ticket, clientId);
        List<Ticket> tickets = new LinkedList<Ticket>();
        while (ticket != null) {
            tickets.add(ticket);
            final long taskId = ticket.getTask().getId();
            sentForExecutionTasksIndex.put(taskId, ticket);
            readyForExecutionTicketsIndex.remove(taskId);

            if (duration >= batchTaskThreshold) {
                break;
            }
            ticket = getTicketFor(clientId);
            if (ticket != null) {
                duration += ticketDurationPredictor.predictDuration(ticket, clientId);
            }
        }
        if (tickets.size() == 1) {
            ticket = tickets.iterator().next();
            LOG.debug("Scheduling one task {} for client #{}", new Object[] {ticket.getTask(), clientId});
            stopWatch.stop("server__scheduler-next-ticket");
            return ticket;
        }

        LOG.debug("Scheduling a batch of tasks: {} for client #{}", new Object[]{tickets, clientId});

        stopWatch.stop("server__scheduler-next-batch-ticket");
        return Ticket.of(new BatchTask(), tickets);
    }

    @Override
    public final void taskFinished(String clientId, long taskId, Object result) {
        final StopWatch stopWatch = StopWatchFactory.createOrGetLocal();
        LOG.debug("Scheduler started handling finished task #{}", new Object[] {taskId});
        // Saving task result
        task2result.put(taskId, result);

        // Checking if task set is finished
        final Long taskSetId = task2taskSet.get(taskId);
        if (taskSetId != null) { // ex. WaitTask has not task set
            final Set<Long> lastTasksIds = taskSet2lastTasks.get(taskSetId);
            if (lastTasksIds != null && task2result.keySet().containsAll(lastTasksIds)) {
                Collection<Object> results = new ArrayList<Object>();
                for (long lastTaskId : lastTasksIds) {
                    results.add(task2result.remove(lastTaskId));
                }
                finishedTaskSets.put(taskSetId, results);
                taskSet2lastTasks.remove(taskSetId);
                long started = taskSet2startTime.remove(taskSetId);
                double time = (System.currentTimeMillis() - started) / 1000.;
                LOG.info("Task set #{} finished in {}s", new Object[] {taskSetId, String.format("%.4f", time)});
            }
        }

        final Ticket finishedTicket = sentForExecutionTasksIndex.remove(taskId);
        stopWatch.lap("server__scheduler-task-finished");

        onTicketFinished(clientId, finishedTicket, result);
        stopWatch.stop("server__scheduler-ontaskfinished");
        LOG.debug("Scheduler finished handling finished task #{}", new Object[] {taskId});
    }

    @Override
    public final boolean isTaskSetFinished(BaseTaskSet baseTaskSet) {
        return finishedTaskSets.containsKey(baseTaskSet.getId());
    }

    @Override
    public final Collection<Object> getTaskSetResults(BaseTaskSet baseTaskSet) {
        return finishedTaskSets.get(baseTaskSet.getId());
    }

    public void setWaitTaskDuration(int waitTaskDuration) {
        this.waitTaskDuration = waitTaskDuration;
    }

    public void setBatchTaskThreshold(int batchTaskThreshold) {
        this.batchTaskThreshold = batchTaskThreshold;
    }

    @Required
    public void setTicketDurationPredictor(TicketDurationPredictor ticketDurationPredictor) {
        this.ticketDurationPredictor = ticketDurationPredictor;
    }

    protected final Collection<Ticket> getReadyForExecutionTickets() {
        return new ArrayList<Ticket>(readyForExecutionTicketsIndex.values());
    }

    protected final void reorganize() {
        final StopWatch stopWatch = StopWatchFactory.createOrGetLocal();
        List<BaseTask> dependingCopy = new ArrayList<BaseTask>(dependingTasksIndex.values());
        for (BaseTask dependingTask : dependingCopy) {
            final List<BaseTask> subTasks = dependingTask.getSubTasks();
            final List<Object> arguments = new ArrayList<Object>(subTasks.size());
            boolean allDone = true;
            for (BaseTask subTask : subTasks) {
                final Object result = task2result.get(subTask.getId());
                if (result == null) {
                    allDone = false;
                    break;
                }
                arguments.add(result);
            }
            if (!allDone) {
                continue;
            }

            dependingTasksIndex.remove(dependingTask.getId());
            readyForExecutionTicketsIndex.put(dependingTask.getId(), Ticket.of(dependingTask, arguments));
        }
        stopWatch.lap("server__scheduler-reorganization");

        onReorganized();
        stopWatch.lap("server__scheduler-onreorganized");
    }

    protected abstract Ticket getTicketFor(String clientId);

    protected void onTicketFinished(String clientId, Ticket ticket, Object result) {
    }

    protected void onReorganized() {
    }

    private static Set<Long> getIds(Collection<BaseTask> tasks) {
        Set<Long> ids = new HashSet<Long>();
        for (BaseTask task : tasks) {
            ids.add(task.getId());
        }
        return ids;
    }
}
